import sqlite3


class SQLiteWrapper:
    def __init__(self, db_name: str, *, cached_statements: int = 100, **kwargs):
        """
        初始化SQLiteWrapper类，创建或连接到指定的SQLite数据库。

        Args:
            db_name (str): 数据库文件名（包括路径，如果不在当前工作目录下）。
            cached_statements (int, optional): 缓存的SQL语句数量，默认为100。
            **kwargs: 其他传递给`sqlite3.connect()`的参数，如`timeout`、`detect_types`等。
        """
        self.db_name = db_name
        self.conn = sqlite3.connect(
            db_name,
            cached_statements=cached_statements,
            **kwargs
        )
        self.conn.row_factory = sqlite3.Row  # 使查询结果以类似字典的方式访问

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        """
        关闭数据库连接。
        """
        if self.conn:
            self.conn.close()
            self.conn = None

    def execute(self, query: str, params=None):
        """
        执行SQL查询或命令。

        Args:
            query (str): SQL查询或命令。
            params (tuple or list, optional): 查询参数，如果没有则为None。

        Returns:
            sqlite3.Cursor: 返回执行查询后得到的游标对象。
        """
        cursor = self.conn.cursor()
        cursor.execute(query, params or ())
        return cursor

    def fetchall(self, query: str, params=None):
        """
        执行SQL查询并获取所有结果。

        Args:
            query (str): SQL查询语句。
            params (tuple or list, optional): 查询参数，如果没有则为None。

        Returns:
            list of sqlite3.Row: 返回查询结果列表，每个元素为一个Row对象，可像字典一样访问列值。
        """
        cursor = self.execute(query, params)
        return cursor.fetchall()

    def fetchone(self, query: str, params=None):
        """
        执行SQL查询并获取第一条结果。

        Args:
            query (str): SQL查询语句。
            params (tuple or list, optional): 查询参数，如果没有则为None。

        Returns:
            sqlite3.Row or None: 返回查询结果的第一条记录，如果没有结果则返回None。
        """
        cursor = self.execute(query, params)
        return cursor.fetchone()

    def commit(self):
        """
        提交当前事务。
        """
        self.conn.commit()

    def rollback(self):
        """
        回滚当前事务。
        """
        self.conn.rollback()

    def create_table(self, table_name: str, columns: str):
        """
        创建一个新表。

        Args:
            table_name (str): 要创建的表名。
            columns (str): 表结构定义，例如 "id INTEGER PRIMARY KEY, name TEXT, age INTEGER"。
        """
        query = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns})"
        self.execute(query).close()
        self.commit()

    def insert_data(self, table_name: str, data: dict):
        """
        插入一条数据到指定表中。

        Args:
            table_name (str): 要插入数据的表名。
            data (dict): 要插入的数据，键为列名，值为对应的值。

        Returns:
            int: 返回插入数据的行ID（主键值），如果表中没有主键，则返回None。
        """
        columns = ", ".join(data.keys())
        placeholders = ", ".join(f":{col}" for col in data.keys())
        query = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        cursor = self.execute(query, tuple(data.values()))
        row_id = cursor.lastrowid if cursor.lastrowid else None
        self.commit()
        return row_id




# 使用示例
with SQLiteWrapper("my_database.db") as db:
    db.create_table("users", "id INTEGER PRIMARY KEY, name TEXT, age INTEGER")
    user_data = {
        "name": "Alice",
        "age": 30,
    }
    inserted_row_id = db.insert_data("users", user_data)

    print(f"Inserted row ID: {inserted_row_id}")
    result = db.fetchone("SELECT * FROM users WHERE id = ?", (1,))
    print(result["name"])
